from large_model.mode_interface.model_interface import ChatInterface
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import MessagesPlaceholder, ChatPromptTemplate
from utils.yaml_utils import save_test_cases
from utils.yaml_template import case_template
from langchain.tools import tool
from large_model.case_template.case_template import TestCase
from langchain.output_parsers import PydanticOutputParser
from typing import List, Dict
import re

# 1. 创建输出解析器
parser = PydanticOutputParser(pydantic_object=TestCase)

# 2. 获取格式指令并转义所有花括号
format_instructions = parser.get_format_instructions()
format_instructions = format_instructions.replace('{', '{{').replace('}', '}}')

# 3. 创建系统提示并转义所有花括号
system_prompt = (
    "你是一个经验丰富的测试专家，需要根据接口文档信息生成接口测试用例。\n\n"
    "请严格遵守以下规则：\n"
    "- 测试用例应为列表嵌套 JSON 格式\n"
    "- 每个测试用例包含断言和对应的参数\n"
    "- 包含正向和反向用例\n"
    "- 覆盖边界值情况\n"
    "- 参数与断言一一对应\n\n"
    "测试用例格式要求：\n"
    f"{format_instructions}\n\n"
    "现在请根据以下接口文档生成测试用例：\n"
)


# 4. 修复工具函数
@tool
def save_test_cases_to_file(test_cases: List[Dict], yaml_filename: str, py_filename: str) -> str:
    """
    保存测试用例到文件系统的工具

    参数:
    test_cases: 生成的测试用例列表（字典格式）
    yaml_filename: 保存的YAML文件名
    py_filename: 生成的测试文件名（不需要.py后缀）

    返回: 执行结果字符串
    """
    try:
        save_test_cases(test_cases, yaml_filename)
        case_template(yaml_filename, py_filename)
        return f"测试用例已保存到 {yaml_filename} 并生成测试文件 {py_filename}.py"
    except Exception as e:
        return f"文件保存失败: {str(e)}"


# 定义工具集
tools = [save_test_cases_to_file]

# 初始化模型
llm = ChatInterface().chat_model()

# 接口文档模板
template = """
info:
  title: 会话详情接口文档
paths:
  url: /trpc.teg_devops.servicedesk.seat.Seat/SessionDetail
  method: get
  headers:
    Content-Type: application/json
    X-Zhiyan-Uid: "v_qxiaochen"
request: params
params:
  proj_id: "项目id，必填， 参数类型为 int, 目前测试环境的项目为 1"
  session_id: "会话id，非必填， 参数类型为 string， （session_id和ticket_id至少填一个）, 推荐使用 1926902512872275968"
  ticket_id: "工单ID，非必填， 参数类型为 int， （session_id和ticket_id至少填一个）"
  web_key: "web站点标识 参数类型为 [string], 非必填"
  source: "来源 参数类型为 [string], 非必填"
"""

# 5. 转义模板中的所有花括号
template = template.replace('{', '{{').replace('}', '}}')
system_message = system_prompt + template

# 6. 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", system_message),
    ("human", "用户请求：{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

# 7. 创建代理
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=3,
    max_execution_time=30,
    handle_parsing_errors=True
)

# 8. 用户请求
user_request = (
    "请根据接口文档生成测试用例，然后使用工具保存结果："
    "yaml_filename=sessiondetail.yaml, "
    "py_filename=test_sessiondetail"
)

# 9. 执行代理
result = agent_executor.invoke({"input": user_request})
print(result["output"])